import os
import unittest

import torch
from torch import hub
from torch.backends import cudnn
from torchsummary import summary

from config import Config
# from training.models import resnet18
from training.models.vit_mlp.vision_transformer import vit_base_patch16_224

CONFIG = Config()
hub.set_dir(CONFIG['TORCH_HOME'])
os.environ["CUDA_VISIBLE_DEVICES"] = CONFIG['CUDA_VISIBLE_DEVICES']

torch.backends.cudnn.benchmark = True


class VisionTransformerTestCase(unittest.TestCase):

    def test_vision_transformer(self):
        self.assertTrue(torch.cuda.is_available())
        # model = vit_small_patch16_224(pretrained=True, num_classes=2)
        model = vit_base_patch16_224(pretrained=True, num_classes=2)
        # model = vit_large_patch16_224(pretrained=True, num_classes=2)

        model = model.cuda()
        input_size = [(3, 224, 224), (3, 224, 224)]
        summary(model, input_size=input_size)


if __name__ == '__main__':
    unittest.main()
